import os

import mindspore
import numpy as np

# 定义数据集路径
data_path = "/tmp/JJK/torch_areas/transformer_learning/text2text/datas"


# 定义数据集
class CmnDataset:
    def __init__(self, data_path: str):
        self.data_path = data_path
        self.en_sents = []
        self.cn_sents = []
        self.load_data()

    def load_data(self):
        cn_path = os.path.join(self.data_path, "cmn_chinese.txt")
        en_path = os.path.join(self.data_path, "cmn_english.txt")
        sign_split = '\t'

        with open(cn_path, "r", encoding="utf-8") as f:
            for line in f:
                self.cn_sents.append(line.strip().split(sign_split))

        with open(en_path, "r", encoding="utf-8") as f:
            for line in f:
                self.en_sents.append(line.strip().split(sign_split))

        assert len(self.en_sents) == len(self.cn_sents), "The number of English and Chinese sentences is not equal."

    def __len__(self):
        return len(self.en_sents)

    def __getitem__(self, idx):
        return self.en_sents[idx], self.cn_sents[idx]


# 实例化数据集
datasets = CmnDataset(data_path)

# 尝试打印数据集中的数据
for en_sent, cn_sent in datasets:
    print("en_sent:", en_sent)
    print("cn_sent:", cn_sent)
    break


# #词典
"""
将每个词元映射到从0开始的数字索引中（为节约存储空间，可过滤掉词频低的词元），词元和数字索引所构成的集合叫做词典（vocabulary）。

以上述“Hello world!”为例，该序列组成的词典为：

{"<unk>": 0, "<pad>": 1, "<bos>": 2, "<eos>": 3, "hello": 4, "world": 5, "!": 6}

在构建词典中，我们使用了4个特殊词元。
<unk>：未知词元（unknown），将出现次数少于一定频率的单词统一判定为未知词元；
<bos>：起始词元（begin of sentence），用来标注一个句子的开始；
<eos>：结束词元（end of sentence），用来标注一个句子的结束；
<pad>：填充词元（padding），当句子长度不够时将句子填充至统一长度；

通过Vocab创建词典后，我们可以实现词元与数字索引之间的互相转换。
我们可以通过调用enocde函数，返回输入词元或者词元序列对应的数字索引或数字索引序列，
反之亦然，我们同样可以通过调用decode函数，返回输入数字索引或数字索引序列对应的词元或词元序列。
"""


class Vocab:
    """通过词频字典，构建词典"""
    special_tokens = ['<unk>', '<pad>', '<bos>', '<eos>']

    def __init__(self, word_count_dict, min_freq=1):
        """
        :param word_count_dict: 字典，key为词元，value为词频
        :param min_freq: 指定需要保留的词频最小值
        """
        self.word2idx = {}
        for idx, tok in enumerate(self.special_tokens):
            self.word2idx[tok] = idx
        filted_dict = {
            w: c
            for w, c in word_count_dict.items() if c >= min_freq
        }
        for w, _ in filted_dict.items():
            self.word2idx[w] = len(self.word2idx)
        self.idx2word = {idx: word for word, idx in self.word2idx.items()}

        self.bos_idx = self.word2idx['<bos>']
        self.eos_idx = self.word2idx['<eos>']
        self.pad_idx = self.word2idx['<pad>']
        self.unk_idx = self.word2idx['<unk>']

    def _word2idx(self, word):
        """单词映射至数字索引"""
        if word not in self.word2idx:
            return self.unk_idx
        return self.word2idx[word]

    def _idx2word(self, idx):
        """数字索引映射至单词"""
        if idx not in self.idx2word:
            raise ValueError('input index is not in vocabulary.')
        return self.idx2word[idx]

    def encode(self, word_or_words):
        """将单个单词或单词数组映射至单个数字索引或数字索引数组"""
        if isinstance(word_or_words, list) or isinstance(word_or_words, np.ndarray):
            return [self._word2idx(i) for i in word_or_words]

        return self._word2idx(word_or_words)

    def decode(self, idx_or_idxs):
        """将单个数字索引或数字索引数组映射至单个单词或单词数组"""
        if isinstance(idx_or_idxs, list) or isinstance(idx_or_idxs, np.ndarray):
            return [self._idx2word(i) for i in idx_or_idxs]

        return self._idx2word(idx_or_idxs)

    def __len__(self):
        return len(self.word2idx)


# Vocab使用举例
word_count_dict = {'hello': 10, 'world': 5, 'good': 3, 'day': 2}
vocab = Vocab(word_count_dict, min_freq=0)
print(vocab.word2idx)
print(vocab.idx2word)
print(vocab.encode('hello'))
print(vocab.decode(2))
print(len(vocab))

# 使用collections中的Counter和OrderedDict统计英/汉语每个单词在整体文本中出现的频率。
# 构建词频字典，然后再将词频字典转为词典。
# 其中，收录所有源语言词元的词典为src_vocab，收录所有目标语言词元的词典为tgt_vocab。

# 在分配数字索引时有一个小技巧：常用的词元对应数值较小的索引(即词频越大的单词的索引越小)，这样可以节约空间。为什么？

# Counter：这是一个用于计数的字典子类，可以快速统计可哈希对象的出现次数。
# OrderedDict：这是一个保留插入顺序的字典，可以高效地按插入顺序保存键值对。
from collections import Counter, OrderedDict


def build_vocab(dataset):
    # 初始化源和目标语言单词的列表，用于存放组成句子的词元（会存放重复的词元，便于词频统计）
    src_words, tgt_words = [], []
    # 对于每对 (src, tgt)（源语句的词元和目标语句的词元）
    for src, tgt in dataset:
        src_words.extend(src)  # 扩展源语言词元列表
        tgt_words.extend(tgt)
    # 使用 Counter 统计 src_words 和 tgt_words 中每个单词的出现频率，并将结果转换为元组列表
    # 即Counter(src_words).items()的结果为 [(词元1, 词元1的词频), (词元2, 词元2的词频), ...]
    src_count_dict = OrderedDict(sorted(Counter(src_words).items(), key=lambda t: t[1], reverse=True))
    tgt_count_dict = OrderedDict(sorted(Counter(tgt_words).items(), key=lambda t: t[1], reverse=True))

    return Vocab(src_count_dict, min_freq=2), Vocab(tgt_count_dict, min_freq=2)


src_vocab, trg_vocab = build_vocab(datasets)
print('Unique tokens in src vocabulary:', len(src_vocab))
print('Unique tokens in tgt vocabulary:', len(trg_vocab))


class Iterator():
    """创建数据迭代器"""

    def __init__(self, dataset, src_vocab, trg_vocab, batch_size, max_len=43, drop_reminder=False):
        """
        :param dataset: 数据集，包含源语言和目标语言的句子对
        :param src_vocab: 源语言词典
        :param trg_vocab: 目标语言词典
        :param batch_size: 批大小
        :param max_len: 最大序列长度
        :param drop_reminder: 是否丢弃最后不满一批的数据。
        """
        self.dataset = dataset
        self.src_vocab = src_vocab
        self.trg_vocab = trg_vocab

        self.batch_size = batch_size
        self.max_len = max_len
        self.drop_reminder = drop_reminder

        length = len(self.dataset) // batch_size
        self.len = length if drop_reminder else length + 1  # 批量数量

    def __call__(self):
        def pad(idx_list, vocab, max_len):
            """统一序列长度，并记录有效长度"""
            idx_pad_list, idx_len = [], []
            for i in idx_list:
                if len(i) > max_len - 2:  # 加上bos和eos后，至少需要max_len-2个词元
                    idx_pad_list.append(
                        [vocab.bos_idx] + i[:max_len - 2] + [vocab.eos_idx]
                    )
                    idx_len.append(max_len)
                else:  # 序列长度小于max_len-2，需要填充pad_idx
                    idx_pad_list.append(
                        [vocab.bos_idx] + i + [vocab.pad_idx] * (max_len - len(i) - 2) + [vocab.eos_idx]
                    )
                    idx_len.append(len(i) + 2)
            return idx_pad_list, idx_len

        def sort_by_length(src, trg):
            """对源语言和目标语言的字段长度进行排序
            具体作用是对输入的源语言（src）和目标语言（trg）序列根据源语言序列的长度进行排序。
            其主要目的是在处理批量数据时，通过将长度相似的序列放在一起，优化模型的训练效率。
            """
            data = zip(src, trg)
            data = sorted(data, key=lambda t: len(t[0]), reverse=True)
            return zip(*list(data))

        def encode_and_pad(batch_data, max_len):
            """将批量中的文本数据转换为数字索引，并统一每个序列的长度"""
            src_data, trg_data = batch_data
            src_idx = [self.src_vocab.encode(i) for i in src_data]
            trg_idx = [self.trg_vocab.encode(i) for i in trg_data]

            src_idx, trg_idx = sort_by_length(src_idx, trg_idx)
            src_idx_pad, src_len = pad(src_idx, src_vocab, max_len)
            trg_idx_pad, _ = pad(trg_idx, trg_vocab, max_len)

            return src_idx_pad, src_len, trg_idx_pad

        for i in range(self.len):
            if i == self.len - 1 and not self.drop_reminder:
                batch_data = self.dataset[i * self.batch_size:]
            else:
                batch_data = self.dataset[i * self.batch_size: (i + 1) * self.batch_size]

            src_idx, src_len, trg_idx = encode_and_pad(batch_data, self.max_len)
            yield mindspore.Tensor(src_idx, mindspore.int32), \
                mindspore.Tensor(src_len, mindspore.int32), \
                mindspore.Tensor(trg_idx, mindspore.int32)

    def __len__(self):
        return self.len


if __name__ == '__main__':

    train_iterator = Iterator(datasets, src_vocab, trg_vocab, batch_size=128, max_len=41, drop_reminder=True)

    for src_idx, src_len, trg_idx in train_iterator():
        print(
            f'src_idx.shape:{src_idx.shape}\n{src_idx}\nsrc_len.shape:{src_len.shape}\n{src_len}\ntrg_idx.shape:{trg_idx.shape}\n{trg_idx}')
        break
